#!/bin/bash
# EbbFlow 本地测试脚本
# 参数: n=6, batchsize=1000, R=5, epoch=10

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 配置参数
N_NODES=6
EPOCHS=10
BATCH_SIZE=1000
R=5
PROJECT_ROOT="/home/ebbflow"
LOG_DIR="${PROJECT_ROOT}/logs_test"
BINARY_NAME="ebbflow_server"
BINARY="${PROJECT_ROOT}/${BINARY_NAME}"

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}==>${NC} ${GREEN}$1${NC}"
}

# 清理函数
cleanup() {
    log_step "清理进程和日志..."
    pkill -f "$BINARY_NAME" 2>/dev/null || true
    sleep 1
    rm -rf "$LOG_DIR"
    mkdir -p "$LOG_DIR"
    log_info "清理完成"
}

# 编译
build() {
    log_step "编译 EbbFlow Server..."
    cd "$PROJECT_ROOT"
    go build -o "$BINARY_NAME" ./src/main/ebbflow_server.go
    if [ $? -ne 0 ]; then
        log_error "编译失败"
        exit 1
    fi
    log_info "编译成功: $BINARY"
}

# 检查配置
check_config() {
    log_step "检查配置文件..."
    
    if [ ! -f "$PROJECT_ROOT/hosts.json" ]; then
        log_error "hosts.json 不存在"
        exit 1
    fi
    
    log_info "配置验证通过: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE, R=$R"
}

# 启动节点
start_nodes() {
    log_step "启动 $N_NODES 个节点..."
    
    for i in $(seq 0 $((N_NODES - 1))); do
        log_info "启动节点 $i..."
        nohup "$BINARY" -id=$i -n=$N_NODES -f=1 -epochs=$EPOCHS -r=$R > "$LOG_DIR/node_${i}.log" 2>&1 &
        echo $! > "$LOG_DIR/node_${i}.pid"
        sleep 0.5
    done
    
    sleep 2
    log_info "所有节点已启动"
}

# 检查节点状态
check_nodes() {
    log_step "检查节点状态..."
    
    local running=0
    for i in $(seq 0 $((N_NODES - 1))); do
        if [ -f "$LOG_DIR/node_${i}.pid" ]; then
            local pid=$(cat "$LOG_DIR/node_${i}.pid")
            if ps -p $pid > /dev/null 2>&1; then
                log_info "节点 $i (PID: $pid) 正在运行"
                running=$((running + 1))
            else
                log_error "节点 $i 已退出"
            fi
        fi
    done
    
    if [ $running -eq $N_NODES ]; then
        log_info "所有 $N_NODES 个节点都在运行"
        return 0
    else
        log_error "只有 $running/$N_NODES 个节点在运行"
        return 1
    fi
}

# 监控测试
monitor_test() {
    log_step "监控测试进度..."
    
    # 计算超时时间（epochs * 30秒每个epoch，加上一些缓冲）
    local timeout=$((EPOCHS * 30 + 60))
    local start_time=$(date +%s)
    
    log_info "监控测试（超时: ${timeout}秒）..."
    log_info "查看实时日志: tail -f $LOG_DIR/node_0.log"
    
    while true; do
        local current_time=$(date +%s)
        local elapsed=$((current_time - start_time))
        
        # 检查超时
        if [ $elapsed -gt $timeout ]; then
            log_error "测试超时（${timeout}秒）"
            return 1
        fi
        
        # 检查是否有节点崩溃
        local running=0
        for i in $(seq 0 $((N_NODES - 1))); do
            if [ -f "$LOG_DIR/node_${i}.pid" ]; then
                local pid=$(cat "$LOG_DIR/node_${i}.pid")
                if ps -p $pid > /dev/null 2>&1; then
                    running=$((running + 1))
                fi
            fi
        done
        
        if [ $running -lt $N_NODES ]; then
            log_error "检测到节点崩溃！运行中: $running/$N_NODES"
            return 1
        fi
        
        # 显示进度
        echo -ne "\r[${elapsed}s] 运行中... (运行节点: $running/$N_NODES)"
        
        sleep 3
    done
}

# 收集结果
collect_results() {
    log_step "收集测试结果..."
    
    echo ""
    echo "================================"
    echo "测试结果摘要"
    echo "================================"
    echo "配置: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE, R=$R"
    echo ""
    
    for i in $(seq 0 $((N_NODES - 1))); do
        echo "--- 节点 $i ---"
        if [ -f "$LOG_DIR/node_${i}.log" ]; then
            # 显示最后几行关键信息
            tail -30 "$LOG_DIR/node_${i}.log" | grep -E "Performance|TPS|Epoch|Completed|Error|error|panic|fatal" || tail -10 "$LOG_DIR/node_${i}.log"
        else
            echo "日志文件不存在"
        fi
        echo ""
    done
    
    # 错误检查
    echo "================================"
    echo "错误检查"
    echo "================================"
    for i in $(seq 0 $((N_NODES - 1))); do
        if [ -f "$LOG_DIR/node_${i}.log" ]; then
            local error_count=$(grep -i "error\|panic\|fatal" "$LOG_DIR/node_${i}.log" | wc -l)
            if [ $error_count -gt 0 ]; then
                echo "节点 $i: $error_count 个错误"
                grep -i "error\|panic\|fatal" "$LOG_DIR/node_${i}.log" | head -5
            else
                echo "节点 $i: 无错误"
            fi
        fi
    done
    echo ""
}

# 主函数
main() {
    echo ""
    echo "================================"
    echo "EbbFlow 本地测试"
    echo "配置: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE, R=$R"
    echo "================================"
    echo ""
    
    # 清理
    cleanup
    
    # 编译
    build
    
    # 检查配置
    check_config
    
    # 启动节点
    start_nodes
    
    # 检查节点状态
    if ! check_nodes; then
        log_error "节点启动失败"
        collect_results
        exit 1
    fi
    
    log_info "所有节点已启动，开始测试..."
    log_info "等待测试完成或按 Ctrl+C 停止..."
    
    # 等待测试完成（根据epochs计算）
    local wait_time=$((EPOCHS * 10 + 30))  # 每个epoch约10秒，加上缓冲
    log_info "预计等待时间: ${wait_time}秒"
    
    sleep $wait_time
    
    # 收集结果
    collect_results
    
    # 停止节点
    log_step "停止所有节点..."
    pkill -f "$BINARY_NAME" 2>/dev/null || true
    sleep 1
    
    log_info "测试完成！"
    echo ""
    log_info "日志文件位置: $LOG_DIR"
    log_info "查看日志: tail -f $LOG_DIR/node_0.log"
}

# 处理 Ctrl+C
trap 'echo ""; log_info "收到中断信号，清理进程..."; pkill -f "$BINARY_NAME" 2>/dev/null || true; exit 130' INT TERM

# 运行
main

